import socket
import json
import random
import struct
import concurrent.futures
import time

SERVER_ADDRESS = ('1.1.1.1', 8888)
KEYWORDS = ["牛逼", "小米", "人性化", "难得", "什么鬼"]
NUM_CONNECTIONS = 100
REQUESTS_PER_CONNECTION = 10
DURATION = 11  # 总持续时间（秒）


def create_connection():
    """建立一个TCP连接并返回连接对象"""
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect(SERVER_ADDRESS)
    return s


def send_requests(connection_id, connection):
    """发送请求并接收响应"""
    print(f"Connection {connection_id} started")
    start_time = time.time()

    for _ in range(DURATION):
        second_start = time.time()
        for _ in range(REQUESTS_PER_CONNECTION):
            keyword = random.choice(KEYWORDS)
            request = {
                "query_id": 1,
                "msg": keyword
            }
            json_request = json.dumps(request) + "\n"

            # 发送请求
            connection.sendall(json_request.encode('utf-8'))

            # 接收响应
            length_bytes = connection.recv(8)
            length = struct.unpack('!Q', length_bytes)[0]

            data = b''
            while len(data) < length:
                chunk = connection.recv(min(4096, length - len(data)))
                if not chunk:
                    raise RuntimeError("socket connection broken")
                data += chunk

            response = json.loads(data.decode('utf-8'))
            print(f"Connection {connection_id} received response: {response}")

        # 等待直到这一秒结束
        time_to_wait = 1 - (time.time() - second_start)
        if time_to_wait > 0:
            time.sleep(time_to_wait)

    connection.close()
    print(f"Connection {connection_id} closed after {time.time() - start_time:.2f} seconds")


def main():
    # 先创建所有连接
    connections = []
    for i in range(NUM_CONNECTIONS):
        conn = create_connection()
        connections.append((i, conn))

    print(f"All {NUM_CONNECTIONS} connections established. Starting requests...")

    # 使用线程池并发发送请求
    with concurrent.futures.ThreadPoolExecutor(max_workers=NUM_CONNECTIONS) as executor:
        futures = [executor.submit(send_requests, conn_id, conn) for conn_id, conn in connections]
        concurrent.futures.wait(futures)


if __name__ == "__main__":
    start_time = time.time()
    main()
    print(f"Total execution time: {time.time() - start_time:.2f} seconds")